#include "yolov3_clip_detector.h"
#include <algorithm>
#include "acl/acl.h"
#include <dirent.h>
#include <math.h>
#include <cstdlib>
#include <chrono>
#include <unistd.h>

Yolov3CLIPDetector::Yolov3CLIPDetector()
{
  std::cout<<"-- Creat instence of Resource Yolov3CLIPDetector"<<std::endl;
}

void Yolov3CLIPDetector::DestroyResource()
{
    done = true;
    queue_input.notify_all();
    queue_pre.notify_all();
    queue_post.notify_all();

    for (auto& th : threads) {
        if (th.joinable()) {
            th.join();
        }
    }

    if (modelDesc != nullptr) {
        (void)aclmdlDestroyDesc(modelDesc);
        modelDesc = nullptr;
    }

    aclError ret = aclmdlUnload(modelId);
    if (ret != ACL_SUCCESS) {
        std::cout<<"Unload model error: "<< ret<<std::endl;
    }

    
    for (size_t i = 0; i < aclmdlGetDatasetNumBuffers(input); ++i) {
        aclDataBuffer* dataBuffer = aclmdlGetDatasetBuffer(input, i);
        aclDestroyDataBuffer(dataBuffer);
        dataBuffer = nullptr;
    }
    aclmdlDestroyDataset(input);
    input = nullptr;

    for (size_t i = 0; i < aclmdlGetDatasetNumBuffers(output); ++i) {
        aclDataBuffer* dataBuffer = aclmdlGetDatasetBuffer(output, i);
        void* data = aclGetDataBufferAddr(dataBuffer);
        (void)aclrtFree(data);
        (void)aclDestroyDataBuffer(dataBuffer);
        dataBuffer = nullptr;
    }

    (void)aclmdlDestroyDataset(output);
    output = nullptr;

    for(uint32_t i=0; i<outinput_nums; i++){
        aclrtFreeHost(outputHostLists[i]);
        aclrtFreeHost(outputDeviceLists[i]);
    }
    
    for(uint32_t i=0; i<input_nums; i++){
        aclrtFree(modelInputs[i]);
    }

    ret = aclrtDestroyContext(context);
    if (ret != ACL_SUCCESS) {
        std::cout<<"destroy context failed, errorCode is : "<< ret<<std::endl;
    }
    context = nullptr;

    ret = aclrtResetDevice(deviceId);
    if (ret != ACL_SUCCESS) {
        std::cout<<"reset device failed, errorCode is : "<< ret<<std::endl;
    }
    std::cout<<"Reset device "<< deviceId<<std::endl; 

    ret = aclFinalize();
    if (ret != ACL_SUCCESS) {
        std::cout<<"finalize acl failed, errorCode is : "<< ret<<std::endl;
    }
    std::cout<<"Finalize acl ok"<<std::endl;
}

void Yolov3CLIPDetector::wait_for_completion() {
    std::unique_lock<std::mutex> lock(results_mutex);
    conda.wait(lock, [this]{ return this->task_counter == 0; });
}

Yolov3CLIPDetector::~Yolov3CLIPDetector()
{
  DestroyResource();
}


void Yolov3CLIPDetector::initThread(){
    
    for (u_int32_t i=0; i<preprocess_thread_num;i++){
        threads.emplace_back(&Yolov3CLIPDetector::modelExecute, this);
    }
    
}

bool Yolov3CLIPDetector::init(const char * ModelPath)
{
    aclError ret = aclInit("");
    if (ret != ACL_SUCCESS) {
        std::cout<<"Acl init failed, errorcode is: "<< ret<<std::endl;
        return false;
    }
    std::cout<<"Acl init ok"<<std::endl;

    // open device
    ret = aclrtSetDevice(deviceId);
    if (ret != ACL_SUCCESS) {
        std::cout<<"Acl open device "<<deviceId<<" failed, errorCode is : "<<ret<<std::endl;
        return false;
    }
    std::cout<<"Open device "<<deviceId<<" ok"<<std::endl;

    ret = aclrtCreateContext(&context, deviceId);
    if (ret != ACL_SUCCESS) {
        std::cout<<"Create acl context failed, errorCode is : "<<ret<<std::endl;
        return false;
    }

    initThread();
    
    // Gets whether the current application is running on host or Device
    ret = aclrtGetRunMode(& g_runMode_);
    if (ret != ACL_SUCCESS) {
        std::cout<<"acl get run mode failed"<<std::endl;
        return false;
    }

    ret = aclmdlLoadFromFile(ModelPath, &modelId);
    if (ret != ACL_SUCCESS) {
        std::cout<<"Load model "<<ModelPath<<" from file return "<<ret<<std::endl;
        return false;
    }

    modelDesc = aclmdlCreateDesc();
    if (modelDesc == nullptr) {
        std::cout<<"create model description failed"<<std::endl;
        return false;
    }

    ret = aclmdlGetDesc(modelDesc, modelId);
    if (ret != ACL_SUCCESS) {
        std::cout<<"get model description failed, modelId is "<<modelId<< ", errorCode is "<<static_cast<int32_t>(ret)<<std::endl;
        return false;
    }

    std::cout<<"create model description success"<<std::endl;
    // cteate output
    output = aclmdlCreateDataset();
    if (output == nullptr) {
        std::cout<<"can't create dataset, create output false"<<std::endl;
        return false;
    }

    outinput_nums = aclmdlGetNumOutputs(modelDesc);
    std::cout<<"outinput_nums: "<<outinput_nums<<std::endl;
    for (size_t i = 0; i < outinput_nums; ++i) {
        const char* name = aclmdlGetOutputNameByIndex(modelDesc, i);
        std::cout<<"output op name: "<<name<<std::endl;
        size_t output_size = aclmdlGetOutputSizeByIndex(modelDesc, i);

        void * outputDevice;
        ret = aclrtMalloc(&outputDevice, output_size, ACL_MEM_MALLOC_NORMAL_ONLY);
        if (ret != ACL_SUCCESS) {
            std::cout<<"can't malloc buffer, size is"<<output_size<<", create output failed"<<std::endl;
            return false;
        }

        aclDataBuffer* outputData = aclCreateDataBuffer(outputDevice, output_size);
        if (ret != ACL_SUCCESS) {
            std::cout<<"can't create data buffer, create output failed"<<std::endl;
            aclrtFree(outputDevice);
            return false;
        }

        ret = aclmdlAddDatasetBuffer(output, outputData);
        if (ret != ACL_SUCCESS) {
            std::cout<<"can't add data buffer, create output failed"<<std::endl;
            aclrtFree(outputDevice);
            aclDestroyDataBuffer(outputData);
            return false;
        }

        void * outputHost;
        ret = aclrtMallocHost(&outputHost, output_size);
        if (ret != ACL_SUCCESS) {
            std::cout<<"Create model input failed"<<std::endl;
            return false;
        }
        outputHostLists.push_back(outputHost);
        outputDeviceLists.push_back(outputDevice);
        outputSizeLists.push_back(output_size);
    }

    std::cout<<"create model output success"<<std::endl;

    std::cout<<"create model INPUT"<<std::endl;
    input = aclmdlCreateDataset();
    if (input == nullptr) {
        std::cout<<"Create input failed for create dataset failed"<<std::endl;
        return false;
    }

    input_nums = aclmdlGetNumInputs(modelDesc);
    std::cout<<"input_nums: "<<input_nums<<std::endl;
    if (input_nums == 0) {
        std::cout<<"Create input failed for no input data"<<std::endl;
        return false;
    }

    for (size_t i= 0; i<input_nums; i++){
        size_t InputSize = aclmdlGetInputSizeByIndex(modelDesc, i);
        std::cout<<"InputSize: "<<InputSize<<std::endl;
        
        void* modelInputBuf;
        aclError aclRet = aclrtMalloc(&modelInputBuf, (size_t)(InputSize), ACL_MEM_MALLOC_HUGE_FIRST);
        if (aclRet != ACL_SUCCESS) {
            std::cout<<"malloc device data buffer failed, aclRet is "<< aclRet<<std::endl;
            return false;
        }

        aclDataBuffer* dataBuf = aclCreateDataBuffer(modelInputBuf, InputSize);
        if (dataBuf == nullptr) {
            std::cout<<"Create data buffer error"<<std::endl;
            return false;
        }

        aclError ret = aclmdlAddDatasetBuffer(input, dataBuf);
        if (ret != ACL_SUCCESS) {
            std::cout<<"Add dataset buffer error "<<ret<<std::endl;
            ret = aclDestroyDataBuffer(dataBuf);
            if (ret != ACL_SUCCESS) {
                std::cout<<"Destroy dataset buffer error "<<ret<<std::endl;
            }
            dataBuf = nullptr;
            return false;
        }
        modelInputs.push_back(modelInputBuf);
        inputSizeLists.push_back(InputSize);
    }
    std::cout<<"create model input success"<<std::endl;
    return true;
}


void Yolov3CLIPDetector::modelExecute()
{
    aclrtSetCurrentContext(context);
    message_pre premsg;
    message_result excutemsg;
    while(!done){
        if (queue_pre.try_pop(premsg)){
            excutemsg.finalOutput.clear();
            std::vector<cv::Mat> frames = premsg.frames;
            std::cout<<"frames size: "<<std::endl;
            {
                std::lock_guard<std::mutex> lock(output_mutex);
                for (size_t i= 0; i<input_nums; i++){
                    std::cout<<"1111"<<std::endl;
                    aclError ret = aclrtMemcpy(modelInputs[i], inputSizeLists[i], frames[i].ptr<uint8_t>(), inputSizeLists[i], ACL_MEMCPY_DEVICE_TO_DEVICE);
                    if (ret != ACL_SUCCESS) {
                        std::cout<<"Copy data to device failed, aclRet is "<<ret<<std::endl;
                        break;
                    }
                }
                
                // Perform reasoning
                std::chrono::high_resolution_clock::time_point t1 = std::chrono::high_resolution_clock::now();
                aclError ret = aclmdlExecute(modelId, input, output);
                std::chrono::high_resolution_clock::time_point t2 = std::chrono::high_resolution_clock::now();
                std::chrono::duration<double, std::milli> time_span = t2 - t1;
                std::cout<<"aclmdlExecute time: "<<time_span.count()<<std::endl;

                if (ret != ACL_SUCCESS) {
                    std::cout<<"execute aclmdlExecute false "<<static_cast<int32_t>(ret)<<std::endl;
                }
                for (size_t i= 0; i<outinput_nums; i++){
                    float * outFloatData0;
                    int outputsize = outputSizeLists[i] /4;
                    if (g_runMode_ == ACL_HOST) {
                        ret = aclrtMemcpy(outputHostLists[i], outputSizeLists[i], outputDeviceLists[i], outputSizeLists[i], ACL_MEMCPY_DEVICE_TO_HOST);
                        outFloatData0 = reinterpret_cast < float * > (outputHostLists[i]);
                    }
                    else
                    {
                        outFloatData0 = reinterpret_cast < float * > (outputDeviceLists[i]);            
                    }
                    std::vector<float> vec_0(outFloatData0, outFloatData0 + outputsize);
                    excutemsg.finalOutput.push_back(vec_0);
                }            
            }
            
            excutemsg.filenames = premsg.filenames;
            std::lock_guard<std::mutex> lock(results_mutex);
            results.push_back(excutemsg);
            if (--task_counter == 0) {
                conda.notify_all();
            }
        }else{
            std::this_thread::sleep_for(std::chrono::milliseconds(10));
        }
    }
}

std::vector<message_result> Yolov3CLIPDetector::inference(std::vector<std::string> &imageLists){
    {
        std::lock_guard<std::mutex> lock(results_mutex);
        results.clear();
        task_counter = 0;
    }


    
    for (int i=0; i<imageLists.size(); i++){
        message_pre premsg;
        cv::Mat img1 = cv::imread(imageLists[i]);
        cv::Mat reiszeMat_visible;
        cv::resize(img1, reiszeMat_visible, cv::Size(g_modelWidth_, g_modelHeight_));
        premsg.filenames.push_back(imageLists[i]);
        premsg.frames.push_back(reiszeMat_visible);
        premsg.filenames.push_back(imageLists[i]);
        premsg.frames.push_back(reiszeMat_visible);
        queue_pre.push(premsg);
        task_counter++;
    }

    wait_for_completion();

    std::vector<message_result> local_results;
    {
        std::lock_guard<std::mutex> lock(results_mutex);
        local_results = results;
    }

    return local_results;

}




